{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "<a target=\"_blank\" href=\"https://colab.research.google.com/github/AI4Finance-Foundation/FinRL-Tutorials/blob/master/2-Advance/Plug_and_Play_with_DRL_Libraries.ipynb\">\n",
    "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
    "</a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "# FinRL-Meta: Plug-and-Play with DRL Libraries\n",
    "Demostration for plug-and-play with ElegantRL, Stable-baselines3, RLlib"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Part 1: Getting Started - Install Python Packages"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    " ### 1.1 Install DRL libraries: FinRL, ElegantRL, RLlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "## install elegantrl library\n",
    "!pip install git+https://github.com/AI4Finance-Foundation/ElegantRL.git\n",
    "## install rllib/ray library\n",
    "!pip install ray[default]\n",
    "## install finrl library\n",
    "!pip install wrds\n",
    "!pip install swig\n",
    "!pip install git+https://github.com/AI4Finance-Foundation/FinRL.git"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### 1.2 Check if the additional packages needed are present, if not install them"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# !pip install trading_calendars\n",
    "# !pip install alpaca_trade_api\n",
    "# !pip install ccxt\n",
    "# !pip install jqdatasdk\n",
    "# !pip install wrds\n",
    "\n",
    "# !pip install lz4\n",
    "# !pip install ray[tune]\n",
    "# !pip install tensorboardX\n",
    "# !pip install gputil"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### 1.3 Import packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "# matplotlib.use('Agg')\n",
    "import datetime\n",
    "\n",
    "import torch\n",
    "import ray\n",
    "from finrl import config\n",
    "from finrl import config_tickers\n",
    "from finrl.meta.preprocessor.yahoodownloader import YahooDownloader\n",
    "from finrl.meta.preprocessor.preprocessors import FeatureEngineer, data_split\n",
    "#from finrl.meta.env_stock_trading.env_stocktrading import StockTradingEnv\n",
    "from finrl.meta.env_stock_trading.env_stocktrading_np import StockTradingEnv\n",
    "from finrl.agents.stablebaselines3.models import DRLAgent as DRLAgent_sb3\n",
    "from finrl.agents.rllib.models import DRLAgent as DRLAgent_rllib\n",
    "from finrl.agents.elegantrl.models import DRLAgent as DRLAgent_erl\n",
    "\n",
    "from finrl.meta.data_processor import DataProcessor\n",
    "from finrl.plot import backtest_stats, backtest_plot, get_daily_return, get_baseline\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Part 2: Train & Test Function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### 2.1 Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def train(start_date, end_date, ticker_list, data_source, time_interval,\n",
    "          technical_indicator_list, drl_lib, env, model_name, if_vix = True,\n",
    "          **kwargs):\n",
    "\n",
    "    #fetch data\n",
    "    DP = DataProcessor(data_source, **kwargs)\n",
    "    df = DP.download_data(ticker_list, start_date, end_date, time_interval)\n",
    "    df = DP.clean_data(df)\n",
    "    df = DP.add_technical_indicator(df,technical_indicator_list)\n",
    "    if if_vix:\n",
    "        df = DP.add_vix(df)\n",
    "    price_array, tech_array, turbulence_array = DP.df_to_array(df, if_vix)\n",
    "    env_config = {'price_array':price_array,\n",
    "              'tech_array':tech_array,\n",
    "              'turbulence_array':turbulence_array,\n",
    "              'if_train':True}\n",
    "    env_instance = env(config=env_config)\n",
    "\n",
    "    #read parameters\n",
    "    cwd = kwargs.get('cwd','./'+str(model_name))\n",
    "\n",
    "    if drl_lib == 'elegantrl':\n",
    "        break_step = kwargs.get('break_step', 1e6)\n",
    "        erl_params = kwargs.get('erl_params')\n",
    "\n",
    "        agent = DRLAgent_erl(env = env,\n",
    "                             price_array = price_array,\n",
    "                             tech_array=tech_array,\n",
    "                             turbulence_array=turbulence_array)\n",
    "\n",
    "        model = agent.get_model(model_name, model_kwargs = erl_params)\n",
    "        trained_model = agent.train_model(model=model,\n",
    "                                          cwd=cwd,\n",
    "                                          total_timesteps=break_step)\n",
    "\n",
    "    elif drl_lib == 'rllib':\n",
    "        total_episodes = kwargs.get('total_episodes', 100)\n",
    "        rllib_params = kwargs.get('rllib_params')\n",
    "\n",
    "        agent_rllib = DRLAgent_rllib(env = env,\n",
    "                       price_array=price_array,\n",
    "                       tech_array=tech_array,\n",
    "                       turbulence_array=turbulence_array)\n",
    "\n",
    "        model,model_config = agent_rllib.get_model(model_name)\n",
    "\n",
    "        model_config['lr'] = rllib_params['lr']\n",
    "        model_config['train_batch_size'] = rllib_params['train_batch_size']\n",
    "        model_config['gamma'] = rllib_params['gamma']\n",
    "\n",
    "        #ray.shutdown()\n",
    "        trained_model = agent_rllib.train_model(model=model,\n",
    "                                          model_name=model_name,\n",
    "                                          model_config=model_config,\n",
    "                                          total_episodes=total_episodes)\n",
    "        trained_model.save(cwd)\n",
    "\n",
    "\n",
    "    elif drl_lib == 'stable_baselines3':\n",
    "        total_timesteps = kwargs.get('total_timesteps', 1e6)\n",
    "        agent_params = kwargs.get('agent_params')\n",
    "\n",
    "        agent = DRLAgent_sb3(env = env_instance)\n",
    "\n",
    "        model = agent.get_model(model_name, model_kwargs = agent_params)\n",
    "        trained_model = agent.train_model(model=model,\n",
    "                                tb_log_name=model_name,\n",
    "                                total_timesteps=total_timesteps)\n",
    "        print('Training finished!')\n",
    "        trained_model.save(cwd)\n",
    "        print('Trained model saved in ' + str(cwd))\n",
    "    else:\n",
    "        raise ValueError('DRL library input is NOT supported. Please check.')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### 2.2 Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def test(start_date, end_date, ticker_list, data_source, time_interval,\n",
    "         technical_indicator_list, drl_lib, env, model_name, if_vix = True,\n",
    "         **kwargs):\n",
    "    #fetch data\n",
    "    DP = DataProcessor(data_source, **kwargs)\n",
    "    df = DP.download_data(ticker_list, start_date, end_date, time_interval)\n",
    "    df = DP.clean_data(df)\n",
    "    df = DP.add_technical_indicator(df, technical_indicator_list)\n",
    "\n",
    "    if if_vix:\n",
    "        df = DP.add_vix(df)\n",
    "    price_array, tech_array, turbulence_array = DP.df_to_array(df, if_vix)\n",
    "\n",
    "    env_config = {'price_array':price_array,\n",
    "            'tech_array':tech_array,\n",
    "            'turbulence_array':turbulence_array,\n",
    "            'if_train':False}\n",
    "    env_instance = env(config=env_config)\n",
    "\n",
    "    #load elegantrl needs state dim, action dim and net dim\n",
    "    net_dimension = kwargs.get('net_dimension', 2**7)\n",
    "    cwd = kwargs.get('cwd','./'+str(model_name))\n",
    "    print(\"price_array: \",len(price_array))\n",
    "\n",
    "    if drl_lib == 'elegantrl':\n",
    "        episode_total_assets = DRLAgent_erl.DRL_prediction(model_name=model_name,\n",
    "                                            cwd=cwd,\n",
    "                                            net_dimension=net_dimension,\n",
    "                                            environment=env_instance)\n",
    "\n",
    "        return episode_total_assets\n",
    "\n",
    "    elif drl_lib == 'rllib':\n",
    "        #load agent\n",
    "        episode_total_assets = DRLAgent_rllib.DRL_prediction(\n",
    "                                  model_name=model_name,\n",
    "                                  env = env,\n",
    "                                  price_array=price_array,\n",
    "                                  tech_array=tech_array,\n",
    "                                  turbulence_array=turbulence_array,\n",
    "                                  agent_path = cwd)\n",
    "\n",
    "        return episode_total_assets\n",
    "\n",
    "\n",
    "    elif drl_lib == 'stable_baselines3':\n",
    "        episode_total_assets = DRLAgent_sb3.DRL_prediction_load_from_file(\n",
    "                                  model_name=model_name,\n",
    "                                  environment = env_instance,\n",
    "                                  cwd = cwd)\n",
    "\n",
    "        return episode_total_assets\n",
    "    else:\n",
    "        raise ValueError('DRL library input is NOT supported. Please check.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Part 3: Set DRL Environment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### 3.1 Get the stock trading env from neo_finrl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from finrl.meta.env_stock_trading.env_stocktrading_np import StockTradingEnv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "import gym\n",
    "from numpy import random as rd\n",
    "\n",
    "class StockTradingEnv(gym.Env):\n",
    "\n",
    "    def __init__(self, config, initial_account=1e6,\n",
    "                 gamma=0.99, turbulence_thresh=99, min_stock_rate=0.1,\n",
    "                 max_stock=1e2, initial_capital=1e6, buy_cost_pct=1e-3,\n",
    "                 sell_cost_pct=1e-3,reward_scaling=2 ** -11,  initial_stocks=None,\n",
    "                 ):\n",
    "        price_ary = config['price_array']\n",
    "        tech_ary = config['tech_array']\n",
    "        turbulence_ary = config['turbulence_array']\n",
    "        if_train = config['if_train']\n",
    "        n = price_ary.shape[0]\n",
    "        self.price_ary =  price_ary.astype(np.float32)\n",
    "        self.tech_ary = tech_ary.astype(np.float32)\n",
    "        self.turbulence_ary = turbulence_ary\n",
    "\n",
    "        self.tech_ary = self.tech_ary * 2 ** -7\n",
    "        self.turbulence_bool = (turbulence_ary > turbulence_thresh).astype(np.float32)\n",
    "        self.turbulence_ary = (self.sigmoid_sign(turbulence_ary, turbulence_thresh) * 2 ** -5).astype(np.float32)\n",
    "\n",
    "        stock_dim = self.price_ary.shape[1]\n",
    "        self.gamma = gamma\n",
    "        self.max_stock = max_stock\n",
    "        self.min_stock_rate = min_stock_rate\n",
    "        self.buy_cost_pct = buy_cost_pct\n",
    "        self.sell_cost_pct = sell_cost_pct\n",
    "        self.reward_scaling = reward_scaling\n",
    "        self.initial_capital = initial_capital\n",
    "        self.initial_stocks = np.zeros(stock_dim, dtype=np.float32) if initial_stocks is None else initial_stocks\n",
    "\n",
    "        # reset()\n",
    "        self.day = None\n",
    "        self.amount = None\n",
    "        self.stocks = None\n",
    "        self.total_asset = None\n",
    "        self.gamma_reward = None\n",
    "        self.initial_total_asset = None\n",
    "\n",
    "        # environment information\n",
    "        self.env_name = 'StockEnv'\n",
    "        # self.state_dim = 1 + 2 + 2 * stock_dim + self.tech_ary.shape[1]\n",
    "        # # amount + (turbulence, turbulence_bool) + (price, stock) * stock_dim + tech_dim\n",
    "        self.state_dim = 1 + 2 + 3 * stock_dim + self.tech_ary.shape[1]\n",
    "        # amount + (turbulence, turbulence_bool) + (price, stock) * stock_dim + tech_dim\n",
    "        self.stocks_cd = None\n",
    "        self.action_dim = stock_dim\n",
    "        self.max_step = self.price_ary.shape[0] - 1\n",
    "        self.if_train = if_train\n",
    "        self.if_discrete = False\n",
    "        self.target_return = 5.0\n",
    "        self.episode_return = 0.0\n",
    "\n",
    "        self.observation_space = gym.spaces.Box(low=-3000, high=3000, shape=(self.state_dim,), dtype=np.float32)\n",
    "        self.action_space = gym.spaces.Box(low=-1, high=1, shape=(self.action_dim,), dtype=np.float32)\n",
    "\n",
    "    def reset(self):\n",
    "        self.day = 0\n",
    "        price = self.price_ary[self.day]\n",
    "\n",
    "        if self.if_train:\n",
    "            self.stocks = (self.initial_stocks + rd.randint(0, 64, size=self.initial_stocks.shape)).astype(np.float32)\n",
    "            self.stocks_cd = np.zeros_like(self.stocks)\n",
    "            self.amount = self.initial_capital * rd.uniform(0.95, 1.05) - (self.stocks * price).sum()\n",
    "        else:\n",
    "            self.stocks = self.initial_stocks.astype(np.float32)\n",
    "            self.stocks_cd = np.zeros_like(self.stocks)\n",
    "            self.amount = self.initial_capital\n",
    "\n",
    "        self.total_asset = self.amount + (self.stocks * price).sum()\n",
    "        self.initial_total_asset = self.total_asset\n",
    "        self.gamma_reward = 0.0\n",
    "        return self.get_state(price)  # state\n",
    "\n",
    "    def step(self, actions):\n",
    "        actions = (actions * self.max_stock).astype(int)\n",
    "\n",
    "        self.day += 1\n",
    "        price = self.price_ary[self.day]\n",
    "        self.stocks_cd += 1\n",
    "\n",
    "        if self.turbulence_bool[self.day] == 0:\n",
    "            min_action = int(self.max_stock * self.min_stock_rate)  # stock_cd\n",
    "            for index in np.where(actions < -min_action)[0]:  # sell_index:\n",
    "                if price[index] > 0:  # Sell only if current asset is > 0\n",
    "                    sell_num_shares = min(self.stocks[index], -actions[index])\n",
    "                    self.stocks[index] -= sell_num_shares\n",
    "                    self.amount += price[index] * sell_num_shares * (1 - self.sell_cost_pct)\n",
    "                    self.stocks_cd[index] = 0\n",
    "            for index in np.where(actions > min_action)[0]:  # buy_index:\n",
    "                if price[index] > 0:  # Buy only if the price is > 0 (no missing data in this particular date)\n",
    "                    buy_num_shares = min(self.amount // price[index], actions[index])\n",
    "                    self.stocks[index] += buy_num_shares\n",
    "                    self.amount -= price[index] * buy_num_shares * (1 + self.buy_cost_pct)\n",
    "                    self.stocks_cd[index] = 0\n",
    "\n",
    "        else:  # sell all when turbulence\n",
    "            self.amount += (self.stocks * price).sum() * (1 - self.sell_cost_pct)\n",
    "            self.stocks[:] = 0\n",
    "            self.stocks_cd[:] = 0\n",
    "\n",
    "        state = self.get_state(price)\n",
    "        total_asset = self.amount + (self.stocks * price).sum()\n",
    "        reward = (total_asset - self.total_asset) * self.reward_scaling\n",
    "        self.total_asset = total_asset\n",
    "\n",
    "        self.gamma_reward = self.gamma_reward * self.gamma + reward\n",
    "        done = self.day == self.max_step\n",
    "        if done:\n",
    "            reward = self.gamma_reward\n",
    "            self.episode_return = total_asset / self.initial_total_asset\n",
    "\n",
    "        return state, reward, done, dict()\n",
    "\n",
    "    def get_state(self, price):\n",
    "        amount = np.array(max(self.amount, 1e4) * (2 ** -12), dtype=np.float32)\n",
    "        scale = np.array(2 ** -6, dtype=np.float32)\n",
    "        return np.hstack((amount,\n",
    "                          self.turbulence_ary[self.day],\n",
    "                          self.turbulence_bool[self.day],\n",
    "                          price * scale,\n",
    "                          self.stocks * scale,\n",
    "                          self.stocks_cd,\n",
    "                          self.tech_ary[self.day],\n",
    "                          ))  # state.astype(np.float32)\n",
    "\n",
    "    @staticmethod\n",
    "    def sigmoid_sign(ary, thresh):\n",
    "        def sigmoid(x):\n",
    "            return 1 / (1 + np.exp(-x * np.e)) - 0.5\n",
    "\n",
    "        return sigmoid(ary / thresh) * thresh\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### 3.2 Set some basic parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "env = StockTradingEnv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "TRAIN_START_DATE = '2014-01-01'\n",
    "TRAIN_END_DATE = '2021-02-28'\n",
    "\n",
    "TEST_START_DATE = '2021-03-01'\n",
    "TEST_END_DATE = '2022-05-01'\n",
    "\n",
    "INDICATORS = ['macd',\n",
    " 'boll_ub',\n",
    " 'boll_lb',\n",
    " 'rsi_30',\n",
    " 'dx_30',\n",
    " 'close_30_sma',\n",
    " 'close_60_sma']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Part 4: Compare the three agents"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### 4.1 eRL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ERL_PARAMS = {\"learning_rate\": 3e-5,\"batch_size\": 2048,\"gamma\":  0.985,\n",
    "        \"seed\":312,\"net_dimension\":512, \"target_step\":5000, \"eval_gap\":60,\n",
    "        \"eval_times\":1}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "#### Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#demo for elegantrl\n",
    "train(start_date = TRAIN_START_DATE,\n",
    "      end_date = TRAIN_END_DATE,\n",
    "      ticker_list = config_tickers.DOW_30_TICKER,\n",
    "      data_source = 'yahoofinance',\n",
    "      time_interval= '1D',\n",
    "      technical_indicator_list= INDICATORS,\n",
    "      drl_lib='elegantrl',\n",
    "      env=env,\n",
    "      model_name='ppo',\n",
    "      cwd='./test_ppo',\n",
    "      erl_params=ERL_PARAMS,\n",
    "      break_step=1e5\n",
    "      )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "#### Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "account_value_erl=test(start_date = TEST_START_DATE,\n",
    "                        end_date = TEST_END_DATE,\n",
    "                        ticker_list = config_tickers.DOW_30_TICKER,\n",
    "                        data_source = 'yahoofinance',\n",
    "                        time_interval= '1D',\n",
    "                        technical_indicator_list= INDICATORS,\n",
    "                        drl_lib='elegantrl',\n",
    "                        env=env,\n",
    "                        model_name='ppo',\n",
    "                        cwd='./test_ppo',\n",
    "                        net_dimension = 512)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "len(account_value_erl)\n",
    "np.save('/content/stock_trading_account_value_erl.npy', account_value_erl)\n",
    "account_value_erl.to_csv('/content/stock_trading_account_value_erl.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "#### Plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "TEST_END_DATE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "DP = DataProcessor('yahoofinance')\n",
    "DP.download_data(ticker_list = [\"^DJI\"],\n",
    "                start_date = TEST_START_DATE,\n",
    "                end_date = TEST_END_DATE,\n",
    "                time_interval = \"1D\")\n",
    "stats = backtest_stats(DP.dataframe, value_col_name = 'close')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "account_value_erl = pd.DataFrame({'date': df.date,'account_value':account_value_erl[0:len(account_value_erl)]})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "account_value_erl.tail()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print(\"==============Get Backtest Results===========\")\n",
    "now = datetime.datetime.now().strftime('%Y%m%d-%Hh%M')\n",
    "\n",
    "perf_stats_all = backtest_stats(account_value=account_value_erl)\n",
    "perf_stats_all = pd.DataFrame(perf_stats_all)\n",
    "perf_stats_all.to_csv(\"./\"+\"/perf_stats_all_\"+now+'.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print(\"==============Compare to DJIA===========\")\n",
    "%matplotlib inline\n",
    "# S&P 500: ^GSPC\n",
    "# Dow Jones Index: ^DJI\n",
    "# NASDAQ 100: ^NDX\n",
    "backtest_plot(account_value_erl,\n",
    "             baseline_ticker = '^DJI',\n",
    "             baseline_start = account_value_erl.loc[0,'date'],\n",
    "             baseline_end = account_value_erl.loc[len(account_value_erl)-1,'date'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### 4.2 RLlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "RLlib_PARAMS = {\"lr\": 5e-6,\"train_batch_size\": 1000,\"gamma\": 0.99}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "#### Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#demo for rllib\n",
    "ray.shutdown() #always shutdown previous session if any\n",
    "\n",
    "train(start_date = TRAIN_START_DATE,\n",
    "      end_date = TRAIN_END_DATE,\n",
    "      ticker_list = config.DOW_30_TICKER,\n",
    "      data_source = 'yahoofinance',\n",
    "      time_interval= '1D',\n",
    "      technical_indicator_list= INDICATORS,\n",
    "      drl_lib='rllib',\n",
    "      env=env,\n",
    "      model_name='ppo',\n",
    "      cwd='./test_ppo',\n",
    "      rllib_params = RLlib_PARAMS,\n",
    "      total_episodes=30)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "#### Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ray.shutdown() #always shutdown previous session if any\n",
    "\n",
    "account_value_rllib = test(start_date = TEST_START_DATE,\n",
    "     end_date = TEST_END_DATE,\n",
    "     ticker_list = config.DOW_30_TICKER,\n",
    "     data_source = 'yahoofinance',\n",
    "     time_interval= '1D',\n",
    "     technical_indicator_list= INDICATORS,\n",
    "     drl_lib='rllib',\n",
    "     env=env,\n",
    "     model_name='ppo',\n",
    "     cwd='./test_ppo/checkpoint_000030/checkpoint-30',\n",
    "     rllib_params = RLlib_PARAMS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "len(account_value_rllib)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "#### Plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "DP = DataProcessor('yahoofinance')\n",
    "DP.download_data(ticker_list = [\"^DJI\"],\n",
    "                start_date = TEST_START_DATE,\n",
    "                end_date = TEST_END_DATE,\n",
    "                time_interval = \"1D\")\n",
    "stats = backtest_stats(DP.dataframe, value_col_name = 'close')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "len(DP.dataframe.date)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "account_value_rllib = pd.DataFrame({'date':DP.dataframe.date,'account_value':account_value_rllib[0:len(account_value_rllib)-1]})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "perf_stats_all = backtest_stats(account_value=account_value_rllib)\n",
    "perf_stats_all = pd.DataFrame(perf_stats_all)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print(\"==============Compare to DJIA===========\")\n",
    "%matplotlib inline\n",
    "# S&P 500: ^GSPC\n",
    "# Dow Jones Index: ^DJI\n",
    "# NASDAQ 100: ^NDX\n",
    "backtest_plot(account_value_rllib,\n",
    "             baseline_ticker = '^DJI',\n",
    "             baseline_start = account_value_rllib.loc[0,'date'],\n",
    "             baseline_end = account_value_rllib.loc[len(account_value_rllib)-1,'date'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### 4.3 Stable-baselines3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "SAC_PARAMS = {\"batch_size\": 128,\"buffer_size\": 100000,\"learning_rate\": 0.0001,\"learning_starts\": 100,\"ent_coef\": \"auto_0.1\",}\n",
    "PPO_PARAMS = {\"n_steps\": 2048,\"ent_coef\": 0.01,\"learning_rate\": 0.00025,\"batch_size\": 128}\n",
    "TD3_PARAMS = {\"batch_size\": 100, \"buffer_size\": 1000000, \"learning_rate\": 0.001}\n",
    "DDPG_PARAMS = {\"batch_size\": 128, \"buffer_size\": 50000, \"learning_rate\": 0.001}\n",
    "A2C_PARAMS = {\"n_steps\": 5, \"ent_coef\": 0.01, \"learning_rate\": 3e-5}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "#### Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#demo for stable-baselines3\n",
    "train(start_date = TRAIN_START_DATE,\n",
    "      end_date = TRAIN_END_DATE,\n",
    "      ticker_list = config.DOW_30_TICKER,\n",
    "      data_source = 'yahoofinance',\n",
    "      time_interval= '1D',\n",
    "      technical_indicator_list= INDICATORS,\n",
    "      drl_lib='stable_baselines3',\n",
    "      env=env,\n",
    "      model_name='sac',\n",
    "      cwd='./test_sac',\n",
    "      agent_params = SAC_PARAMS,\n",
    "      total_timesteps=1e4)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "#### Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "account_value_sb3=test(start_date = TEST_START_DATE,\n",
    "     end_date = TEST_END_DATE,\n",
    "     ticker_list = config.DOW_30_TICKER,\n",
    "     data_source = 'yahoofinance',\n",
    "     time_interval= '1D',\n",
    "     technical_indicator_list= INDICATORS,\n",
    "     drl_lib='stable_baselines3',\n",
    "     env=env,\n",
    "     model_name='sac',\n",
    "     cwd='./test_sac.zip')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "len(account_value_sb3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "#### Plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "DP = DataProcessor('yahoofinance')\n",
    "DP.download_data(ticker_list = [\"^DJI\"],\n",
    "                start_date = TEST_START_DATE,\n",
    "                end_date = TEST_END_DATE,\n",
    "                time_interval = \"1D\")\n",
    "stats = backtest_stats(DP.dataframe, value_col_name = 'close')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "account_value_sb3 = pd.DataFrame({'date':DP.dataframe.date,'account_value':account_value_sb3[0:len(account_value_sb3)-1]})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "perf_stats_all = backtest_stats(account_value=account_value_sb3)\n",
    "perf_stats_all = pd.DataFrame(perf_stats_all)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "account_value_sb3.tail()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print(\"==============Compare to DJIA===========\")\n",
    "#%matplotlib inline\n",
    "# S&P 500: ^GSPC\n",
    "# Dow Jones Index: ^DJI\n",
    "# NASDAQ 100: ^NDX\n",
    "backtest_plot(account_value_sb3,\n",
    "             baseline_ticker = '^DJI',\n",
    "             baseline_start = account_value_sb3.loc[0,'date'],\n",
    "             baseline_end = account_value_sb3.loc[len(account_value_sb3)-1,'date'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Part 5: Use Plotly to compare eRL, RLlib and SB3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "DP.dataframe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from datetime import datetime as dt\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import plotly\n",
    "import plotly.graph_objs as go"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "daily_return = account_value_sb3.copy()\n",
    "daily_return['sb3_return'] = account_value_sb3.account_value.pct_change()\n",
    "daily_return['erl_return'] = account_value_erl.account_value.pct_change()\n",
    "daily_return['rllib_return'] = account_value_rllib.account_value.pct_change()\n",
    "daily_return['djia_return'] = DP.dataframe.adjcp.pct_change()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "daily_return.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "daily_return.to_csv('daily_return_erl_sb3_rllib.csv',index=False)\n",
    "#daily_return = pd.read_csv('daily_return_erl_sb3_rllib.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "rllib_cumpod =(daily_return.rllib_return+1).cumprod()-1\n",
    "sb3_cumpod =(daily_return.sb3_return+1).cumprod()-1\n",
    "erl_cumpod =(daily_return.erl_return+1).cumprod()-1\n",
    "dji_cumpod =(daily_return.djia_return+1).cumprod()-1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "time_ind = pd.Series(daily_return.date)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "trace0_portfolio = go.Scatter(x = time_ind, y = rllib_cumpod, mode = 'lines', name = 'RLlib')\n",
    "\n",
    "trace1_portfolio = go.Scatter(x = time_ind, y = dji_cumpod, mode = 'lines', name = 'DJIA')\n",
    "trace2_portfolio = go.Scatter(x = time_ind, y = sb3_cumpod, mode = 'lines', name = 'Stablebaselines3')\n",
    "trace3_portfolio = go.Scatter(x = time_ind, y = erl_cumpod, mode = 'lines', name = 'ElegantRL')\n",
    "#trace4_portfolio = go.Scatter(x = time_ind, y = addpg_cumpod, mode = 'lines', name = 'Adaptive-DDPG')\n",
    "#trace5_portfolio = go.Scatter(x = time_ind, y = min_cumpod, mode = 'lines', name = 'Min-Variance')\n",
    "\n",
    "#trace4 = go.Scatter(x = time_ind, y = addpg_cumpod, mode = 'lines', name = 'Adaptive-DDPG')\n",
    "\n",
    "#trace2 = go.Scatter(x = time_ind, y = portfolio_cost_minv, mode = 'lines', name = 'Min-Variance')\n",
    "#trace3 = go.Scatter(x = time_ind, y = spx_value, mode = 'lines', name = 'SPX')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig = go.Figure()\n",
    "fig.add_trace(trace3_portfolio)\n",
    "fig.add_trace(trace2_portfolio)\n",
    "\n",
    "fig.add_trace(trace0_portfolio)\n",
    "fig.add_trace(trace1_portfolio)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "fig.update_layout(\n",
    "    legend=dict(\n",
    "        x=0,\n",
    "        y=1,\n",
    "        traceorder=\"normal\",\n",
    "        font=dict(\n",
    "            family=\"sans-serif\",\n",
    "            size=15,\n",
    "            color=\"black\"\n",
    "        ),\n",
    "        bgcolor=\"White\",\n",
    "        bordercolor=\"white\",\n",
    "        borderwidth=2\n",
    "\n",
    "    ),\n",
    ")\n",
    "#fig.update_layout(legend_orientation=\"h\")\n",
    "fig.update_layout(title={\n",
    "        #'text': \"Cumulative Return using FinRL\",\n",
    "        'y':0.85,\n",
    "        'x':0.5,\n",
    "        'xanchor': 'center',\n",
    "        'yanchor': 'top'})\n",
    "#with Transaction cost\n",
    "#fig.update_layout(title =  'Quarterly Trade Date')\n",
    "fig.update_layout(\n",
    "#    margin=dict(l=20, r=20, t=20, b=20),\n",
    "\n",
    "    paper_bgcolor='rgba(1,1,0,0)',\n",
    "    plot_bgcolor='rgba(1, 1, 0, 0)',\n",
    "    #xaxis_title=\"Date\",\n",
    "    yaxis_title=\"Cumulative Return\",\n",
    "xaxis={'type': 'date',\n",
    "       'tick0': time_ind[0],\n",
    "        'tickmode': 'linear',\n",
    "       'dtick': 86400000.0 *70}\n",
    "\n",
    ")\n",
    "fig.update_xaxes(showline=True,linecolor='black',showgrid=True, gridwidth=1, gridcolor='LightSteelBlue',mirror=True)\n",
    "fig.update_yaxes(showline=True,linecolor='black',showgrid=True, gridwidth=1, gridcolor='LightSteelBlue',mirror=True)\n",
    "fig.update_yaxes(zeroline=True, zerolinewidth=1, zerolinecolor='LightSteelBlue')\n",
    "\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "Demo_Plug_and_Play_with_DRL_Libraries.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
